﻿#pragma once
using namespace std;
#include <iostream>
#include <vector>
#include <assert.h>

enum colour
{
	BLACK,
	RED
};

template<class K, class V>
struct RBTNode
{
	pair<K, V> _kv;
	RBTNode* _left;
	RBTNode* _right;
	RBTNode* _parent;
	colour _col;

	RBTNode(const pair<K, V>& kv)
		:_kv(kv)
		, _left(nullptr)
		, _right(nullptr)
		, _parent(nullptr)
	{
	}
};

template<class K, class V>
struct RBTNode
{
	pair<K, V> _kv;
	RBTNode<K, V>* _left;
	RBTNode<K, V>* _right;
	RBTNode<K, V>* _parent;
	colour _col;

	RBTNode(const pair<K, V>& kv)
		:_kv(kv)
		, _left(nullptr)
		, _right(nullptr)
		, _parent(nullptr)
	{
	}
};

template<class K, class V>
class RBTree
{
	typedef RBTNode<K, V> Node;
public:
	bool Insert(const pair<K, V> kv)
	{
		if (_root == nullptr)
		{
			_root = new Node(kv);
			_root->_col = BLACK; //根结点为黑
			return true;
		}
		Node* parent = nullptr;
		Node* cur = _root;
		while (cur)
		{
			if (cur->_kv.first > kv.first)
			{
				parent = cur;
				cur = cur->_left;
			}
			else if (cur->_kv.first < kv.first)
			{
				parent = cur;
				cur = cur->_right;
			}
			else
			{
				return false;
			}
		}
		cur = new Node(kv);
		cur->_col = RED;   //插入结点默认为红
		if (kv.first < parent->_kv.first)
			parent->_left = cur;
		else
			parent->_right = cur;
		cur->_parent = parent;

		while (parent && parent->_col == RED) //更新到根结点和parent为黑出循环停止更新
		{
			Node* grandparent = parent->_parent;
			if (parent == grandparent->_left)  //p在左，u在右
			{
				Node* uncle = grandparent->_right;
				if (uncle && uncle->_col == RED)     //u存在且为红
				{
					parent->_col = uncle->_col = BLACK;
					grandparent->_col = RED;
					//继续向上调整
					cur = grandparent;
					parent = cur->_parent;
				}
				else                                //u存在且为黑或不存在
				{
					if (cur == parent->_left) //都左边小，右单旋
					{
						//      g
						//    p   u
						//  c
						RotateR(grandparent);
						grandparent->_col = RED;
						parent->_col = BLACK;
					}
					else                     //左边小的右边大，左右双旋
					{
						//      g
						//    p   u
						//     c
						RotateL(parent);
						RotateR(grandparent);
						grandparent->_col = RED;
						cur->_col = BLACK; //cur被推上来做根
					}
					break; //情况2、3执行完后便更新完毕
				}
			}
			else  //p在右，u在左
			{
				Node* uncle = grandparent->_left;
				if (uncle && uncle->_col == RED)     //u存在且为红
				{
					parent->_col = uncle->_col = BLACK;
					grandparent->_col = RED;
					//继续向上调整
					cur = grandparent;
					parent = cur->_parent;
				}
				else                                //u存在且为黑或不存在
				{
					if (cur == parent->_right) //都右边小，左单旋
					{
						//      g
						//    u   p
						//         c
						RotateL(grandparent);
						grandparent->_col = RED;
						parent->_col = BLACK;
					}
					else                     //右边大的左边小，右左双旋
					{
						//      g
						//    u   p
						//       c
						RotateR(parent);
						RotateL(grandparent);
						grandparent->_col = RED;
						cur->_col = BLACK; //cur被推上来做根
					}
					break;
				}
			}
		}
		_root->_col = BLACK;  //情况1cur更新到根节点需要把根结点变黑
		return true;
	}

	Node* Find(const K& key)
	{
		Node* cur = _root;
		while (cur)
		{
			if (cur->_kv.first < key)
			{
				cur = cur->_right;
			}
			else if (cur->_kv.first > key)
			{
				cur = cur->_left;
			}
			else
			{
				return cur;
			}
		}
		return nullptr;
	}

	void InOrder()
	{
		_InOrder(_root);
		cout << endl;
	}

	bool IsBalance()
	{
		if (_root == nullptr)
		{
			return true;  //空树也是红黑树
		}
		if (_root->_col == RED)
		{
			return false; // 违反规则二
		}

		Node* cur = _root;
		int blacknum = 0;
		while (cur != nullptr) //用其中一条路径的黑色结点个数作为参考值
		{
			if (cur->_col == BLACK)
				++blacknum;
			cur = cur->_left;
		}
		return	Check(_root, blacknum, 0);
	}

	//递归前序遍历检查
	bool Check(Node* root, int blacknum, int curnum)
	{
		if (root == nullptr)
		{
			if (blacknum == curnum)
				return true;
			else
				cout << "存在黑色结点的数量不相等的路径" << endl;
				return false; //路径黑色结点个数不匹配，违反规则四
		}
		//当前结点和父节点都为红，违反规则三
		if (root->_col == RED && root->_parent && root->_parent->_col == RED)
		{
			cout << root->_kv.first << "存在连续的红⾊结点" << endl;
			return false;
		}
		if (root->_col == BLACK)
		{
			++curnum;
		}
		return Check(root->_left, blacknum, curnum)
			&& Check(root->_right, blacknum, curnum);
	}

	void _InOrder(Node* root)
	{
		if (root == nullptr)
			return;
		_InOrder(root->_left);
		cout << root->_kv.first << " ";
		_InOrder(root->_right);
	}

	void RotateR(Node* parent)
	{
		Node* subL = parent->_left;
		Node* subLR = subL->_right;
		//分配孩子
		subL->_right = parent;
		parent->_left = subLR;
		//调整subLR的父亲
		if (subLR)
			subLR->_parent = parent;
		Node* parentParent = parent->_parent; //提前保存
		//调整parent的父亲
		parent->_parent = subL;
		//调整subL的父亲和parentParent的孩子
		if (parent == _root)
		{
			_root = subL; 
			subL->_parent = nullptr;
		}
		else
		{
			subL->_parent = parentParent;
			if (parent == parentParent->_left)
				parentParent->_left = subL;
			else
				parentParent->_right = subL;
		}
	}

	void RotateL(Node* parent)
	{
		Node* subR = parent->_right;
		Node* subRL = subR->_left;
		//分配孩子
		subR->_left = parent;
		parent->_right = subRL;
		//调整subLR的父亲
		if (subRL)
			subRL->_parent = parent;
		Node* parentParent = parent->_parent; //提前保存
		//调整parent的父亲
		parent->_parent = subR;
		//调整subL的父亲和parentParent的孩子
		if (parent == _root)
		{
			_root = subR;
			subR->_parent = nullptr;
		}
		else
		{
			subR->_parent = parentParent;
			if (parent == parentParent->_left)
				parentParent->_left = subR;
			else
				parentParent->_right = subR;
		}
	}

private:
	Node* _root = nullptr;
};



